#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/6/15 9:27
# @Author  : LiShan
# @Email   : lishan_1997@126.com
# @File    : main.py
# @Note    : this is note
import os

import cv2.cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

save_path = "./MOG_FG_BG"
width = 480
height = 360

if os.path.exists(save_path):
    pass
else:
    os.mkdir(save_path)


def point_distance_line(point,line_point1,line_point2):
    #计算向量
    vec1 = line_point1 - point
    vec2 = line_point2 - point
    distance = np.abs(np.cross(vec1,vec2)) / np.linalg.norm(line_point1-line_point2)
    return distance

# def extend_line(x1, y1, x2, y2, w, h):
#     # print(666)

def max_line(L1, L2):
    # L1：直线1，L2：直线2
    x1, y1, x2, y2 = L1
    x3, y3, x4, y4 = L2
    # 取最高的起点
    if y1 < y3:
        x5, y5 = x1, y1
    else:
        x5, y5 = x3, y3
    # 取最低的终点
    if y2 > y4:
        x6, y6 = x2, y2
    else:
        x6, y6 = x4, y4
    # 返回合并后的直线
    return x5, y5, x6, y6

class TrafficFlowDetected:
    def __init__(self):
        self.video = "./video/test.mp4"
        self.width = 640
        self.height = 480
        self.out_w = self.width
        self.out_h = self.height

    # # 车道线检测
    # def lane_line_detect(self, img):
    #     """
    #         for i in range(len(lane_line)):
    #             x1, y1, x2, y2 = lane_line[i]
    #             cv.line(oshow, (x1, y1), (x2, y2), (0, 0, 255), 2)
    #     """
    #     width = self.width
    #     height = self.height
    #     img = cv.resize(img, (width, height))
    #     # 原图
    #     origin = img
    #     # 转为灰度图
    #     gray = cv.cvtColor(origin, cv.COLOR_BGR2GRAY)
    #     edges = cv.Canny(gray, 50, 150, apertureSize=3)
    #     # 转为RGB图
    #     orgb = cv.cvtColor(origin, cv.COLOR_BGR2RGB)
    #     oshow = orgb.copy()
    #     # 霍夫概率变换寻找直线，image、rho、theta、threshold、minLineLength、maxLineGap
    #     lines = cv.HoughLinesP(edges, 1, np.pi / 180, 60, minLineLength=100, maxLineGap=80)
    #     lane_line = []
    #     lane_k = []
    #     for line in lines:
    #         x1, y1, x2, y2 = line[0]
    #         # 求直线斜率
    #         if x2 != x1:
    #             k = (y2 - y1) / (x2 - x1)
    #         else:
    #             k = (y2 - y1) / (x2 - x1 + np.finfo(float).eps)
    #         # 排除较为水平的直线
    #         if np.abs(k) > 0.45:
    #             for i in range(len(lane_k)):
    #                 if np.abs(k - lane_k[i]) < 0.1:
    #                     # 合并相近直线为最长直线
    #                     L1 = lane_line[i]
    #                     L2 = line[0]
    #                     p0 = np.array([x1, y1])
    #                     p1 = np.array([L1[0], L1[1]])
    #                     p2 = np.array([L1[2], L1[3]])
    #                     d = point_distance_line(p0, p1, p2)
    #                     if d < 3:
    #                         x1, y1, x2, y2 = max_line(L1, L2)
    #                         lane_line[i] = x1, y1, x2, y2
    #                         break
    #             else:
    #                 lane_k.append(k)
    #                 lane_line.append([x1, y1, x2, y2])
    #     return lane_line

    # 判断目标框位置，删除内部框
    def is_inside(self, o, i):
        # o为第一个框，i为第二个框
        ox, oy, ow, oh = o
        ix, iy, iw, ih = i
        # 如果第一个矩形框被完全包含在第二个矩形框中，可确定第一个矩形框应该被丢弃
        return ox > ix and oy > iy and ox + ow < ix + iw and oy + oh < iy + ih

    # 车流量检测
    def count_detected(self, video=None):
        if video is None:
            video = self.video
        count = 0

        # 路径、编码器、帧率、尺寸、是否彩色
        cap = cv.VideoCapture(video)
        fps = cap.get(cv.CAP_PROP_FPS)
        # print(fps)
        save_name = video[:-4] + ".avi"
        # 常用编码器：'DIVX'(.avi)、'MJPG'(.mp4)、'XVID'(.avi,推荐)、'X264'(.mkv)
        fourcc = cv.VideoWriter_fourcc(*'XVID')
        out_size = (self.out_w, self.out_h)
        out = cv.VideoWriter(save_name, fourcc, 25, out_size, True)
        # 窗口排列
        # cv.namedWindow("origin", 0)
        # cv.resizeWindow("origin", self.width, self.height)
        # cv.moveWindow("origin", 0, 0)

        cv.namedWindow("gray", 0)
        cv.resizeWindow("gray", self.width, self.height)
        cv.moveWindow("gray", self.width * 1, 0)

        cv.namedWindow("fgmask", 0)
        cv.resizeWindow("fgmask", self.width, self.height)
        cv.moveWindow("fgmask", self.width * 2, 0)

        cv.namedWindow("median", 0)
        cv.resizeWindow("median", self.width, self.height)
        cv.moveWindow("median", self.width * 0, 30 + self.height * 1)

        cv.namedWindow("morphology", 0)
        cv.resizeWindow("morphology", self.width, self.height)
        cv.moveWindow("morphology", self.width * 1, 30 + self.height * 1)

        cv.namedWindow("Contours", 0)
        cv.resizeWindow("Contours", self.width, self.height)
        cv.moveWindow("Contours", self.width * 2, 30 + self.height * 1)

        # 背景差分法（混合高斯建模）创建一个对象
        # detectShadows表示是否检测阴影，默认为True，会降低速度
        subtractor = cv.createBackgroundSubtractorMOG2(detectShadows=False)
        frame_num = 0
        try:
            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    pass
                else:
                    frame_num += 1
                    # shape:高度(行)、宽度(列)、通道数
                    self.out_h = frame.shape[0]
                    self.out_w = frame.shape[1]
                    out_size = (self.out_w, self.out_h)
                    out = cv.VideoWriter(save_name, fourcc, 25, out_size, True)
                    # 调整图片大小：原图、（宽，高）
                    frame = cv.resize(frame, (self.width, self.height))
                    # 原图
                    origin = frame
                    # cv.imshow("origin", origin)
                    # 转为灰度图
                    gray = cv.cvtColor(origin, cv.COLOR_BGR2GRAY)
                    # cv.imshow("gray", gray)
                    # ret,Binary = cv.threshold(gary,50,255,cv.THRESH_BINARY)

                    # 车道线检测
                    # lane_line = self.lane_line_detect(frame)
                    # for i in range(len(lane_line)):
                    #     x1, y1, x2, y2 = lane_line[i]
                    #     cv.line(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
                    # 车流量检测线
                    dis = 40
                    start_point = (0, self.height - dis)
                    end_point = (self.width, self.height - dis)
                    # cv.line(frame, start_point, end_point, (0, 255, 0), 2, 4)
                    # cv.line(frame, start_point, end_point, (0, 255, 0), 2, 4)
                    # 检测器虚拟线圈
                    x = int(self.width * 0.1 / 2)
                    y = self.height - 100
                    w = int(self.width * 0.9)
                    h = 80
                    VC = frame[x:x+w, y:y+h]
                    cv.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
                    # 根据背景差分法提取前景腌膜
                    fgmask = subtractor.apply(gray)
                    bgmask = subtractor.getBackgroundImage(gray)
                    cv.imwrite(save_path + "/" + str(frame_num) + "_fgmask" +   ".jpg",fgmask)
                    cv.imwrite(save_path + "/" + str(frame_num) + "_bgmask" +   ".jpg",bgmask)
                    # cv.imshow("fgmask", fgmask)
                    # 中值滤波
                    median = cv.medianBlur(fgmask, 5)
                    # cv.imshow("median", median)
                    # 形态学操作
                    element = cv.getStructuringElement(cv.MORPH_RECT, (1, 1))
                    element2 = cv.getStructuringElement(cv.MORPH_RECT, (1, 1))
                    # 开运算
                    image = cv.morphologyEx(median, cv.MORPH_OPEN, element)
                    # 膨胀运算
                    morphology = cv.dilate(image, element2)
                    # cv.imshow('morphology', morphology)
                    # 二值图像提取轮廓，轮廓跟踪算法（Suzuki，1985）
                    contours, hierarchy = cv.findContours(morphology, cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
                    pre_object = []
                    for cnt in contours:
                        # 依次取出每一条轮廓，计算点集或灰度图像的非零像素的右上边界矩形
                        x, y, w, h = cv.boundingRect(cnt)
                        rect = x, y, w, h
                        # if w + h >= self.width:
                        #     count += 1
                        area = self.width * self.height
                        # 过滤掉小物体、大物体
                        if cv.contourArea(cnt) < area * 0.01 or w < self.width * 0.05 or h < self.height * 0.05:
                            continue
                        else:
                            pre_object.append(rect)
                    # print("预选目标框个数：%d" % len(pre_object))
                    # 遍历检测结果来丢弃不含有检测目标的区域
                    object = []
                    num = 0
                    # ri、qi为框的编号，r、q为对应框的x，y，w，h坐标参数
                    for ri, r in enumerate(pre_object):
                        for qi, q in enumerate(pre_object):
                            if ri != qi and self.is_inside(r, q):
                                break
                            else:
                                num += 1
                        if num == len(pre_object):
                            object.append(r)
                        num = 0
                    # print("目标框个数：%d" % len(object))
                    for i in range(len(object)):
                        x, y, w, h = object[i]
                        cv.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 1)
                        x_c = int(x + w / 2)
                        y_c = int(y + h / 2)
                        if y_c > self.height-dis:
                            count += 1
                    cv.putText(frame, "fps: " + str(int(fps)), (self.width - 120, 30), cv.LINE_AA, 0.8,
                               (0, 255, 0), 2)
                    cv.putText(frame, "flow: " + str(count), (self.width - 120, 30 + 30), cv.LINE_AA, 0.8,
                               (0, 255, 0), 2)
                    # cv.imshow("Contours", frame)

                    hist = cv.calcHist([VC], [0, 1], None, [180, 256], [0, 180, 0, 256])
                    # plt.imshow(hist, interpolation='nearest')
                    # cv.imshow("hist", hist)

                    # plt.hist(VC.ravel(), 256, [0, 256])
                    # plt.show()

                    frame = cv.resize(frame, (self.out_w, self.out_h))
                    out.write(frame)
                    k = cv.waitKey(1) & 0xff
                    if k == 27:
                        break
            cap.release()
            cv.destroyAllWindows()
        except:
            pass


if __name__ == '__main__':
    TrafficFlowDetected().count_detected("./video/background_model.mp4")
